import torch
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from torchvision.datasets import MNIST
import mlflow.pytorch

'''
此段代码主要用于解释如何使用`mlflow.pytorch.load_model`
'''

# load and normalize the dataset
transform = transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.5,), (0.5,)),
                              ])

df_test = MNIST("dataset", download=True, train=False, transform=transform)
test_dataloader = DataLoader(df_test, batch_size=32, shuffle=True)
model_uri = 'runs:/42926a185dd34b6abfec60e1f411c972/model'
loaded_model = mlflow.pytorch.load_model(model_uri)
predicted=[]
gt=[]
acc = 0
correct_num = 0
print(len(test_dataloader))
with torch.no_grad():
    n_correct=0
    n_samples=0
    for images,labels in test_dataloader:
        images=images.reshape(-1,784)
        output=loaded_model(images) #applying the model we have built
        labels=labels
        _,prediction=torch.max(output,1)
        res = sum(x == y for x, y in zip(prediction.tolist(), labels.tolist()))
        acc += res/len(labels)

#print(predicted)
#print(gt)
print("accuracy: ",acc/len(test_dataloader))